import gym
from gym import spaces
import numpy as np

class SnakeEnv(gym.Env):
    def __init__(self, width=10, height=10):
        super(SnakeEnv, self).__init__()

        # 定义游戏区域的大小
        self.width = width
        self.height = height

        # 定义动作空间
        self.action_space = spaces.Discrete(4)  # 0: 上, 1: 下, 2: 左, 3: 右

        # 定义状态空间
        self.observation_space = spaces.Box(low=0, high=2, shape=(self.width, self.height), dtype=np.uint8)

        # 初始化游戏状态
        self.reset()

    def reset(self):
        # 初始化游戏区域
        self.game_area = np.zeros((self.height, self.width), dtype=np.uint8)

        # 初始化贪吃蛇的位置和长度
        self.snake_position = [(self.width // 2, self.height // 2)]
        self.snake_length = 1

        # 在游戏区域中表示贪吃蛇
        for position in self.snake_position:
            self.game_area[position] = 1

        # 初始化食物的位置
        self.food_position = self._generate_food_position()
        self.game_area[self.food_position] = 2

        return self.game_area

    def step(self, action):
        # 获取当前贪吃蛇的头部位置
        head_x, head_y = self.snake_position[0]

        # 根据动作更新贪吃蛇的头部位置
        if action == 0:  # 上
            head_y -= 1
        elif action == 1:  # 下
            head_y += 1
        elif action == 2:  # 左
            head_x -= 1
        elif action == 3:  # 右
            head_x += 1

        # 检查贪吃蛇是否撞到边界或自己的身体
        if head_x < 0 or head_x >= self.width or head_y < 0 or head_y >= self.height or (head_x, head_y) in self.snake_position[1:]:
            done = True
            reward = -10
        else:
            done = False
            reward = 0

            # 更新贪吃蛇的位置
            self.snake_position.insert(0, (head_x, head_y))
            self.game_area[head_x, head_y] = 1

            # 检查贪吃蛇是否吃到食物
            if (head_x, head_y) == self.food_position:
                self.snake_length += 1
                reward = 10
                self.food_position = self._generate_food_position()
                self.game_area[self.food_position] = 2
            else:
                # 如果没有吃到食物，移除贪吃蛇的尾部
                tail_x, tail_y = self.snake_position.pop()
                self.game_area[tail_x, tail_y] = 0

        return self.game_area, reward, done, {}

    def render(self, mode='human'):
        if mode == 'human':
            print(self.game_area)
        else:
            raise NotImplementedError

    def _generate_food_position(self):
        # 生成一个随机的食物位置，确保不在贪吃蛇的身体上
        while True:
            food_x = np.random.randint(0, self.width)
            food_y = np.random.randint(0, self.height)
            if (food_x, food_y) not in self.snake_position:
                return food_x, food_y
